# coding: utf-8
import numpy as np


class SGD:
    """随机梯度下降法（Stochastic Gradient Descent）"""

    def __init__(self, lr=0.01):
        self.lr = lr

    def update(self, update_iter):
        for var, grad in update_iter:
            var -= self.lr * grad


class Adam:
    """Adam (http://arxiv.org/abs/1412.6980v8)"""

    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.iter = 0
        self.m = None
        self.v = None

    def update(self, update_iter):
        if self.m is None:
            self.m, self.v = [], []
            idx = 0
            for var, grad in update_iter:
                self.m.append(np.zeros_like(var))
                self.v.append(np.zeros_like(var))
                idx += 1

        self.iter += 1
        lr_t = self.lr * np.sqrt(1.0 - self.beta2 ** self.iter) / (1.0 - self.beta1 ** self.iter)

        idx = 0
        for var, grad in update_iter:
            self.m[idx] += (1 - self.beta1) * (grad - self.m[idx])
            self.v[idx] += (1 - self.beta2) * (grad ** 2 - self.v[idx])

            var -= lr_t * self.m[idx] / (np.sqrt(self.v[idx]) + 1e-7)
            idx += 1
